import numpy as np

from ..utils import TransformationUtils


class RangeBearingSensor:
    def __init__(self, noise: np.ndarray) -> None:
        super().__init__()

        self._noise = noise

    def H(self, xv, xi, enable_noise=False):
        z = np.array([
            np.linalg.norm(xi[:2] - xv[:2]),
            TransformationUtils.wrap(np.arctan2(xi[1] - xv[1], xi[0] - xv[0]) - xv[2])[0]
        ])
        if enable_noise:
            z += np.random.randn(2) * self._noise
        return z

    def Hxv(self, xv, xi):
        r = np.linalg.norm(xv[:2] - xi[:2])
        r2 = r * r

        hxv = np.array([
            [(xv[0] - xi[0]) / r, (xv[1] - xi[1]) / r, 0],
            [(xi[1] - xv[1]) / r2, (xv[0] - xi[0]) / r2, -1]
        ])

        return hxv

    def Hxi(self, xv, xi):
        r = np.linalg.norm(xv[:2] - xi[:2])
        r2 = r * r

        hxi = np.array([
            [(xi[0] - xv[0]) / r, (xi[1] - xv[1]) / r],
            [(xv[1] - xi[1]) / r2, (xi[0] - xv[0]) / r2]
        ])

        return hxi

    def Hw(self, xv, xi):
        return np.eye(2)

    def G(self, xv, z):
        g = np.array([
            xv[0] + z[0] * np.cos(xv[2] + z[1]),
            xv[1] + z[0] * np.sin(xv[2] + z[1])
        ])
        return g

    def Gxi(self, xv, z):
        gxi = np.array([
            [1, 0, -z[0] * np.sin(xv[2] + z[1])],
            [0, 1, z[0] * np.cos(xv[2] + z[1])]
        ])
        return gxi

    def Gz(self, xv, z):
        gz = np.array([
            [np.cos(xv[2] + z[1]), -z[0] * np.sin(xv[2] + z[1])],
            [np.sin(xv[2] + z[1]), z[0] * np.cos(xv[2] + z[1])]
        ])
        return gz

    def handle_measure(self, z: np.ndarray):
        for i in range(1, z.size, 2):
            z[i] = TransformationUtils.wrap(z[i])[0]
        return z
